package cn.longwingstech.intelligence.agent.app.memory;

import cn.longwingstech.intelligence.agent.model.entity.App;
import cn.longwingstech.intelligence.agent.model.entity.ChatHistory;
import cn.longwingstech.intelligence.agent.service.AppService;
import cn.longwingstech.intelligence.agent.service.ChatHistoryService;
import com.alibaba.fastjson2.JSON;
import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.memory.ChatMemoryRepository;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.SystemMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.stereotype.Component;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.ConcurrentHashMap;

@Slf4j
@Component
public class RedisSpringAiChatMemory implements ChatMemoryRepository {
    @Resource
    private StringRedisTemplate stringRedisTemplate;
    private final String REDIS_MEMORY_KEY = "chat_memory_key:";
    @Resource
    private ChatHistoryService chatHistoryService;
    @Resource
    private AppService appService;
    
    // 添加缓存来避免重复查询数据库
    private final ConcurrentHashMap<String, List<Message>> messageCache = new ConcurrentHashMap<>();
    private final ConcurrentHashMap<String, Long> cacheTimestamp = new ConcurrentHashMap<>();
    private static final long CACHE_EXPIRE_TIME = 5 * 60 * 1000; // 5分钟缓存过期时间

    @Override
    public List<String> findConversationIds() {
        return List.of();
    }

    @Override
    public List<Message> findByConversationId(String conversationId) {
        try {
            // 输入验证
            if (conversationId == null || conversationId.trim().isEmpty()) {
                log.warn("conversationId 为空，返回空消息列表");
                return new ArrayList<>();
            }
            
            // 检查缓存
            if (isCacheValid(conversationId)) {
                log.debug("从缓存中获取消息，conversationId: {}", conversationId);
                return new ArrayList<>(messageCache.get(conversationId));
            }
            
            log.debug("从Redis获取消息，conversationId: {}", conversationId);
            String json = stringRedisTemplate.opsForValue().get(REDIS_MEMORY_KEY + conversationId);
            List<Message> messages = new ArrayList<>();
            
            if (json == null) {
                // 从数据库加载历史消息
                messages = loadMessagesFromDatabase(conversationId);
            } else {
                // 从Redis加载消息
                messages = loadMessagesFromRedis(json);
            }
            
            // 限制消息数量，避免内存溢出
            if (messages.size() > 20) {
                messages = messages.subList(Math.max(0, messages.size() - 20), messages.size());
                log.info("消息数量超过限制，已截取最近20条消息，conversationId: {}", conversationId);
            }
            
            // 更新缓存
            updateCache(conversationId, messages);
            
            return new ArrayList<>(messages);
            
        } catch (Exception e) {
            log.error("获取对话消息失败，conversationId: {}", conversationId, e);
            return new ArrayList<>();
        }
    }
    
    private boolean isCacheValid(String conversationId) {
        Long timestamp = cacheTimestamp.get(conversationId);
        if (timestamp == null) {
            return false;
        }
        return System.currentTimeMillis() - timestamp < CACHE_EXPIRE_TIME;
    }
    
    private void updateCache(String conversationId, List<Message> messages) {
        try {
            messageCache.put(conversationId, new ArrayList<>(messages));
            cacheTimestamp.put(conversationId, System.currentTimeMillis());
            
            // 清理过期缓存
            cleanExpiredCache();
        } catch (Exception e) {
            log.error("更新缓存失败，conversationId: {}", conversationId, e);
        }
    }
    
    private void cleanExpiredCache() {
        try {
            long currentTime = System.currentTimeMillis();
            cacheTimestamp.entrySet().removeIf(entry -> {
                if (currentTime - entry.getValue() > CACHE_EXPIRE_TIME) {
                    messageCache.remove(entry.getKey());
                    return true;
                }
                return false;
            });
        } catch (Exception e) {
            log.error("清理过期缓存失败", e);
        }
    }
    
    private List<Message> loadMessagesFromDatabase(String conversationId) {
        try {
            App app = appService.getOne(new LambdaQueryWrapper<App>().eq(App::getChatId, conversationId));
            if (app == null) {
                log.debug("未找到对应的App，conversationId: {}", conversationId);
                return new ArrayList<>();
            }
            
            Long appid = app.getId();
            List<ChatHistory> historyList = chatHistoryService.list(
                new LambdaQueryWrapper<ChatHistory>()
                    .eq(ChatHistory::getAppid, appid)
                    .orderByAsc(ChatHistory::getCreateTime) // 确保按时间顺序
                    .last("LIMIT 20") // 限制查询数量
            );
            
            List<Message> messages = new ArrayList<>();
            for (ChatHistory history : historyList) {
                try {
                    String messageType = history.getMessageType();
                    String messageContent = history.getMessage();
                    
                    if (messageContent == null || messageContent.trim().isEmpty()) {
                        continue;
                    }
                    
                    Message message = null;
                    if ("user".equals(messageType)) {
                        message = new UserMessage(messageContent);
                    } else if ("assistant".equals(messageType)) {
                        message = new AssistantMessage(messageContent);
                    } else if ("system".equals(messageType)) {
                        message = new SystemMessage(messageContent);
                    }
                    
                    if (message != null) {
                        messages.add(message);
                    }
                } catch (Exception e) {
                    log.error("处理历史消息失败，historyId: {}", history.getId(), e);
                }
            }
            
            // 移除最后一条消息（如果存在）
            if (!messages.isEmpty()) {
                messages.remove(messages.size() - 1);
            }
            
            log.debug("从数据库加载了{}条消息，conversationId: {}", messages.size(), conversationId);
            return messages;
            
        } catch (Exception e) {
            log.error("从数据库加载消息失败，conversationId: {}", conversationId, e);
            return new ArrayList<>();
        }
    }
    
    private List<Message> loadMessagesFromRedis(String json) {
        try {
            List<ChatEntity> chatEntities = JSON.parseArray(json, ChatEntity.class);
            List<Message> messages = new ArrayList<>();
            
            for (ChatEntity chatEntity : chatEntities) {
                try {
                    String type = chatEntity.getType();
                    String text = chatEntity.getText();
                    
                    if (text == null || text.trim().isEmpty()) {
                        continue;
                    }
                    
                    Message message = null;
                    switch (type) {
                        case "USER": {
                            message = new UserMessage(text);
                            break;
                        }
                        case "ASSISTANT": {
                            message = new AssistantMessage(text);
                            break;
                        }
                        case "SYSTEM": {
                            message = new SystemMessage(text);
                            break;
                        }
                        default: {
                            log.debug("未知的消息类型: {}", type);
                        }
                    }
                    
                    if (message != null) {
                        messages.add(message);
                    }
                } catch (Exception e) {
                    log.error("处理Redis消息失败", e);
                }
            }
            
            log.debug("从Redis加载了{}条消息", messages.size());
            return messages;
            
        } catch (Exception e) {
            log.error("从Redis解析消息失败", e);
            return new ArrayList<>();
        }
    }

    @Override
    public void saveAll(String conversationId, List<Message> messages) {
        try {
            // 输入验证
            if (conversationId == null || conversationId.trim().isEmpty()) {
                log.warn("conversationId 为空，跳过保存");
                return;
            }
            
            if (messages == null || messages.isEmpty()) {
                log.debug("消息列表为空，跳过保存，conversationId: {}", conversationId);
                return;
            }
            
            // 限制保存的消息数量
            List<Message> messagesToSave = messages;
            if (messages.size() > 20) {
                messagesToSave = messages.subList(Math.max(0, messages.size() - 20), messages.size());
                log.info("消息数量超过限制，只保存最近20条消息，conversationId: {}", conversationId);
            }
            
            List<ChatEntity> chatEntities = new ArrayList<>();
            for (Message message : messagesToSave) {
                try {
                    ChatEntity chatEntity = new ChatEntity();
                    chatEntity.setType(message.getMessageType().name());
                    chatEntity.setText(message.getText());
                    chatEntities.add(chatEntity);
                } catch (Exception e) {
                    log.error("转换消息失败", e);
                }
            }
            
            String jsonString = JSON.toJSONString(chatEntities);
            stringRedisTemplate.opsForValue().set(
                REDIS_MEMORY_KEY + conversationId, 
                jsonString,
                60 * 60, 
                TimeUnit.SECONDS
            );
            
            // 更新缓存
            updateCache(conversationId, messagesToSave);
            
            log.debug("保存了{}条消息到Redis，conversationId: {}", chatEntities.size(), conversationId);
            
        } catch (Exception e) {
            log.error("保存消息到Redis失败，conversationId: {}", conversationId, e);
        }
    }

    @Override
    public void deleteByConversationId(String conversationId) {
        try {
            if (conversationId == null || conversationId.trim().isEmpty()) {
                log.warn("conversationId 为空，跳过删除");
                return;
            }
            
            stringRedisTemplate.delete(REDIS_MEMORY_KEY + conversationId);
            
            // 清理缓存
            messageCache.remove(conversationId);
            cacheTimestamp.remove(conversationId);
            
            log.debug("删除对话消息，conversationId: {}", conversationId);
            
        } catch (Exception e) {
            log.error("删除对话消息失败，conversationId: {}", conversationId, e);
        }
    }
}